import datetime
import talib
import pandas as pd
import numpy as np
from numpy.lib.stride_tricks import as_strided

stime = '20190503'
etime = '20210805'
# component = "000016.SH"
# index_syn = "000001.SH"
component = "399344.SZ"
index_syn = "399001.SZ"

# 399006.SZ
stime1 = datetime.datetime.strptime(stime, '%Y%m%d').strftime('%Y-%m-%d')
etime1 = datetime.datetime.strptime(etime, '%Y%m%d').strftime('%Y-%m-%d')
stock_list = get_index_stocks(component)

data = get_price(stock_list, stime, etime,
                 '1d',
                 #        close                    涨跌幅        振幅         换手率
                 ['open', 'close', 'high', 'low', 'volume', 'quote_rate', 'amp_rate', 'turnover_rate'],
                 True,  # 是否跳过停牌
                 "pre",  # 前复权
                 0,  # 天数
                 is_panel=1)
df = data.to_frame().reset_index()
df.rename(columns={'major': 'date', 'minor': 'symbol'}, inplace=True)
# print(df.columns)
df = df.sort_values(['symbol', 'date'])
df = df.reset_index(drop=True)
sszs = get_price([index_syn], stime, etime,
                 '1d',
                 #   close
                 ['close'],
                 True,  # 是否跳过停牌
                 "pre",  # 前复权
                 0,  # 天数
                 is_panel=0)[index_syn]
sszs.rename(columns={"close": "szss_close"}, inplace=True)

df = pd.merge(df, sszs, left_on="date", right_index=True)
df["day"] = df["date"].apply(lambda x: x.strftime("%Y-%m-%d"))
# print(df.dtypes)
# print(df.head())
# print(0/0)
# df.to_csv("./stock_with_szzs.csv")
# 周线换手率均值, 月线换手率均值
df['close_ma3'] = df.groupby(['symbol'])['close'].rolling(3).mean().reset_index(drop=True, level=0)
df['close_ma5'] = df.groupby(['symbol'])['close'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_ma10'] = df.groupby(['symbol'])['close'].rolling(10).mean().reset_index(drop=True, level=0)
df['close_ma20'] = df.groupby(['symbol'])['close'].rolling(20).mean().reset_index(drop=True, level=0)
df['close_ma60'] = df.groupby(['symbol'])['close'].rolling(60).mean().reset_index(drop=True, level=0)

# df["close_ma" + "_fib_w1"] = df['close_ma5'] / df['close_ma20'] - 1.0
# df["close_ma"  + "_fib_m1"] = df['close_ma20'] / df['close_ma60'] - 1.0

# 周线换手率均值, 月线换手率均值
df['turnover_r5'] = df.groupby(['symbol'])['turnover_rate'].rolling(5).mean().reset_index(drop=True, level=0)
df['turnover_r20'] = df.groupby(['symbol'])['turnover_rate'].rolling(20).mean().reset_index(drop=True, level=0)
df['turnover_r60'] = df.groupby(['symbol'])['turnover_rate'].rolling(60).mean().reset_index(drop=True, level=0)

df["turnover_ma" + "_fib_w1"] = df['turnover_r5'] / df['turnover_r20'] - 1.0
df["turnover_ma" + "_fib_m1"] = df['turnover_r20'] / df['turnover_r60'] - 1.0

df['return1'] = 1.0 - df['close'] / df.groupby(['symbol'])['close'].shift(-1)
# df['return3'] = 1.0 - df['close'] / df.groupby(['symbol'])['close'].shift(-3)
df['return5'] = 1.0 - df['close'] / df.groupby(['symbol'])['close'].shift(-5)  #


# df['return10'] = 1.0 - df['close'] / df.groupby(['symbol'])['close'].shift(-10)
# df['return15'] = 1.0 - df['close'] / df.groupby(['symbol'])['close'].shift(-15)
# df['return20'] = 1.0 - df['close'] / df.groupby(['symbol'])['close'].shift(-20)
def corr_new(df1, col1, col2, n_list=[5, 10, 15, 20, 60]):
    for n in n_list:
        df1["corr" + str(n)] = df1[col1].astype("float32").rolling(n).corr(df1[col2].astype("float32"))
    return df1


def corr_new2(df1, col1, col2, n_list=[5, 10, 15, 20, 60]):
    for n in n_list:
        df1["corr" + str(n)] = pd.rolling_corr(df1[col1].astype("float32"), df1[col2].astype("float32"), n)
    return df1


def talib_pattern(df1):
    df1['crows2'] = talib.CDL2CROWS(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['crows3'] = talib.CDL3BLACKCROWS(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['inside3'] = talib.CDL3INSIDE(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['line_strike'] = talib.CDL3LINESTRIKE(df1['open'].values, df1['high'].values, df1['low'].values,
                                              df1['close'].values)
    df1['outside'] = talib.CDL3OUTSIDE(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['stars_in_south3'] = talib.CDL3STARSINSOUTH(df1['open'].values, df1['high'].values, df1['low'].values,
                                                    df1['close'].values)
    df1['wihite_soliders3'] = talib.CDL3WHITESOLDIERS(df1['open'].values, df1['high'].values, df1['low'].values,
                                                      df1['close'].values)
    return df1


def talib_mom(df1):
    df1['adx_14'] = talib.ADX(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['adxr_14'] = talib.ADXR(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['apo'] = talib.APO(df1['close'].values, fastperiod=12, slowperiod=26, matype=0)
    df1['aroon_up'], df1["aroon_down"] = talib.AROON(df1['high'].values, df1['low'].values, timeperiod=14)
    df1['aroon_osc'] = talib.AROONOSC(df1['high'].values, df1['low'].values, timeperiod=14)
    df1['bop'] = talib.BOP(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['cci'] = talib.CCI(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['cmo'] = talib.CMO(df1['close'].values, timeperiod=14)

    df1['dx'] = talib.DX(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1["macd"], df1["macd_signal"], df1["macd_hist"] = talib.MACD(df1['close'].values, fastperiod=12, slowperiod=26,
                                                                   signalperiod=9)
    df1['mfi'] = talib.MFI(df1['high'].values, df1['low'].values, df1['close'].values, df1['volume'].values,
                           timeperiod=14)
    df1['di'] = talib.MINUS_DI(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['dm'] = talib.MINUS_DM(df1['high'].values, df1['low'].values, timeperiod=14)
    df1['mom'] = talib.MOM(df1['close'].values, timeperiod=10)

    df1['plus_di'] = talib.PLUS_DI(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['plus_dm'] = talib.PLUS_DM(df1['high'].values, df1['low'].values, timeperiod=14)
    # 。。。待续
    return df1


def calc_var(stock):
    # symbol = df1["symbol"].iloc[0]
    data = get_price([stock], stime, etime,
                     '1m',
                     #   close
                     ['close'],
                     True,  # 是否跳过停牌
                     "pre",  # 前复权
                     0,  # 天数
                     is_panel=1)

    dfm = data.to_frame().reset_index()
    dfm.rename(columns={'major': 'date', 'minor': 'symbol'}, inplace=True)
    dfm["day"] = dfm["date"].apply(lambda x: x.strftime("%Y-%m-%d"))
    dfm["minute"] = dfm["date"].apply(lambda x: x.strftime("%H:%M"))
    dfm = dfm.sort_values(['symbol', 'date'])
    dfm = dfm.reset_index(drop=True)
    ## 需要产生一条日期数据
    ## todo 按照小时计算var
    ## todo 计算和指数的相关系数
    # df_var = dfm.groupby(['symbol', "day"])[['close']].std()
    df_var = dfm.groupby(['symbol', "day"])['close'].agg({"close_var": "std", "close_var_f": lambda x: x[:60].std(),
                                                          "close_var_l": lambda x: x[180:].std()})
    df_var = df_var.reset_index()  ## 将groupby还原
    # print(df_var.head())
    # df_var.rename(columns={'close': 'close_var'}, inplace=True)

    return df_var


def calc_var_hour(df):
    std_day = df['close'].std()
    std_fhour = df[df['minute'] < '10:30']['close'].std()
    std_lhour = df[df['minute'] > '14:00']['close'].std()
    return pd.DataFrame([[std_day, std_fhour, std_lhour]], columns={'close_var', 'close_fhour_var', 'close_lhour_var'},
                        index=[df[['symbol', "day"]][0]])
    # return [std_day,std_fhour,std_lhour]


df = df.round(3)
# print("++++++++++")
# df_var = calc_var()
# print("-------")
# pd.merge(df, df_var, left_on=["symbol", "day"], right_index=True, how="inner")
cache_df = None
ii = 0
print("len:", len(stock_list))
for stock in stock_list:
    print("-------" + str(ii))
    ii = ii + 1
    mdf = pd.merge(df[df['symbol'] == stock], calc_var(stock), on=["symbol", "day"])
    # print(type(mdf.index))
    if cache_df is None:
        cache_df = mdf
    else:
        cache_df = pd.concat([cache_df, mdf])
df = cache_df
# print(df.dtypes)
print(df.head())
# df = df.reset_index()
df = df.sort_values(['symbol', 'date'])
df = df.reset_index(drop=True)

df['close_var_ma3'] = df.groupby(['symbol'])['close_var'].rolling(3).mean().reset_index(drop=True, level=0)
df['close_var_ma5'] = df.groupby(['symbol'])['close_var'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_var_ma10'] = df.groupby(['symbol'])['close_var'].rolling(10).mean().reset_index(drop=True, level=0)
df['close_var_ma20'] = df.groupby(['symbol'])['close_var'].rolling(20).mean().reset_index(drop=True, level=0)
df['close_var_ma60'] = df.groupby(['symbol'])['close_var'].rolling(60).mean().reset_index(drop=True, level=0)
df["close_var_ma" + "_fib_w1"] = df['close_var_ma5'] / df['close_var_ma20'] - 1.0
df["close_var_ma" + "_fib_m1"] = df['close_var_ma20'] / df['close_var_ma60'] - 1.0

df['close_var_f_ma5'] = df.groupby(['symbol'])['close_var_f'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_var_f_ma20'] = df.groupby(['symbol'])['close_var_f'].rolling(20).mean().reset_index(drop=True, level=0)
df['close_var_f_ma60'] = df.groupby(['symbol'])['close_var_f'].rolling(60).mean().reset_index(drop=True, level=0)
df["close_var_f_ma" + "_fib_w1"] = df['close_var_f_ma5'] / df['close_var_f_ma20'] - 1.0
df["close_var_f_ma" + "_fib_m1"] = df['close_var_f_ma20'] / df['close_var_f_ma60'] - 1.0

df['close_var_l_ma5'] = df.groupby(['symbol'])['close_var_l'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_var_l_ma20'] = df.groupby(['symbol'])['close_var_l'].rolling(20).mean().reset_index(drop=True, level=0)
df['close_var_l_ma60'] = df.groupby(['symbol'])['close_var_l'].rolling(60).mean().reset_index(drop=True, level=0)
df["close_var_l_ma" + "_fib_w1"] = df['close_var_l_ma5'] / df['close_var_l_ma20'] - 1.0
df["close_var_l_ma" + "_fib_m1"] = df['close_var_l_ma20'] / df['close_var_l_ma60'] - 1.0


df['close_var_ma3_upd'] =  df['close_var_ma3']  / df.groupby(['symbol'])['close_var_ma3'].shift(3)   - 1.0
df['close_var_ma5_upd'] =  df['close_var_ma5']  / df.groupby(['symbol'])['close_var_ma5'].shift(5)   - 1.0
df['close_var_ma10_upd'] = df['close_var_ma10'] / df.groupby(['symbol'])['close_var_ma10'].shift(10) - 1.0


df['close_ma3_upd'] =   df['close_ma3']  / df.groupby(['symbol'])['close_ma3'].shift(3) - 1.0
df['close_ma5_upd'] =   df['close_ma5']  / df.groupby(['symbol'])['close_ma5'].shift(5) - 1.0
df['close_ma10_upd'] =  df['close_ma10'] / df.groupby(['symbol'])['close_ma10'].shift(10) - 1.0

df['var_sub_close_ma_rate3'] = df['close_var_ma3_upd'] - df['close_ma3_upd']
df['var_sub_close_ma_rate5'] = df['close_var_ma5_upd'] - df['close_ma5_upd']
df['var_sub_close_ma_rate10'] = df['close_var_ma10_upd'] - df['close_ma10_upd']

df = df.round(3)
df.drop(['open', 'high', 'low', 'volume'], axis=1, inplace=True)
print(df.head())
df.sort_index(inplace=True)
df.to_csv("./stock_399344_var_model_data_v2.csv", index=False)